from fastmcp import FastMCP
from pydantic import BaseModel
import paramiko
import asyncio
from typing import Optional
import logging
import json

mcp = FastMCP("Demo 🚀")
# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class SSHConnection(BaseModel):
    """SSH连接信息模型
    
    Attributes:
        hostname (str): SSH服务器主机名或IP地址
        port (int): SSH服务器端口号，默认为22
        username (str): SSH登录用户名
        password (Optional[str]): SSH登录密码，如果使用密钥认证则为None
        key_filename (Optional[str]): SSH私钥文件路径，如果使用密码认证则为None
    """
    hostname: str
    port: int = 22
    username: str
    password: Optional[str] = None
    key_filename: Optional[str] = None

class CommandRequest(BaseModel):
    """命令执行请求模型
    
    Attributes:
        command (str): 要在远程服务器上执行的命令
        connection (SSHConnection): SSH连接信息
    """
    command: str
    connection: SSHConnection

class CommandResponse(BaseModel):
    """命令执行响应模型
    
    Attributes:
        success (bool): 命令是否执行成功
        output (str): 命令的标准输出内容
        error (Optional[str]): 命令的错误输出内容，如果执行成功则为None
    """
    success: bool
    output: str
    error: Optional[str] = None

class SSHClientManager:
    """SSH客户端管理器"""
    
    def __init__(self):
        self.client = None
    
    def connect(self, connection: SSHConnection) -> bool:
        """连接到SSH服务器
        
        Args:
            connection (SSHConnection): SSH连接配置信息
            
        Returns:
            bool: 连接成功返回True，失败返回False
        """
        try:
            self.client = paramiko.SSHClient()
            self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            # 连接参数
            connect_kwargs = {
                'hostname': connection.hostname,
                'port': connection.port,
                'username': connection.username
            }
            
            if connection.password:
                connect_kwargs['password'] = connection.password
            elif connection.key_filename:
                connect_kwargs['key_filename'] = connection.key_filename
            
            self.client.connect(**connect_kwargs)
            logger.info(f"成功连接到 {connection.hostname}:{connection.port}")
            return True
            
        except Exception as e:
            logger.error(f"SSH连接失败: {str(e)}")
            self.client = None
            return False
    
    def execute_command(self, command: str) -> CommandResponse:
        """执行命令
        
        Args:
            command (str): 要在远程服务器上执行的命令
            
        Returns:
            CommandResponse: 命令执行结果，包含成功状态、输出内容和错误信息
        """
        if not self.client:
            return CommandResponse(
                success=False,
                output="",
                error="SSH客户端未连接"
            )
        
        try:
            stdin, stdout, stderr = self.client.exec_command(command)
            
            # 读取输出
            output = stdout.read().decode('utf-8', errors='ignore')
            error_output = stderr.read().decode('utf-8', errors='ignore')
            
            # 检查退出状态
            exit_status = stdout.channel.recv_exit_status()
            
            if exit_status == 0:
                return CommandResponse(
                    success=True,
                    output=output,
                    error=error_output if error_output else None
                )
            else:
                return CommandResponse(
                    success=False,
                    output=output,
                    error=error_output if error_output else f"命令执行失败，退出状态: {exit_status}"
                )
                
        except Exception as e:
            logger.error(f"命令执行失败: {str(e)}")
            return CommandResponse(
                success=False,
                output="",
                error=str(e)
            )
    
    def disconnect(self):
        """断开SSH连接
        
        如果当前有活动的SSH连接，则关闭连接并清理资源
        """
        if self.client:
            self.client.close()
            self.client = None
            logger.info("SSH连接已断开")

# 全局SSH管理器实例
ssh_manager = SSHClientManager()

@mcp.tool
def connect(hostname: str, username: str, password: str = None, port: int = 22, key_filename: str = None) -> dict:
    """连接到SSH服务器
    
    Args:
        hostname (str): SSH服务器主机名或IP地址
        username (str): SSH登录用户名
        password (str, optional): SSH登录密码，如果使用密钥认证则不需要
        port (int, optional): SSH服务器端口号，默认为22
        key_filename (str, optional): SSH私钥文件路径，如果使用密码认证则不需要
        
    Returns:
        dict: 连接结果，包含以下字段：
            - success (bool): 连接是否成功
            - message (str): 连接状态消息
            - connection_info (dict, optional): 连接信息（仅在成功时包含）
    """
    try:
        # 创建连接对象
        connection = SSHConnection(
            hostname=hostname,
            port=port,
            username=username,
            password=password,
            key_filename=key_filename
        )
        
        success = ssh_manager.connect(connection)
        if success:
            return {
                "success": True,
                "message": f"成功连接到 {hostname}:{port}",
                "connection_info": {
                    "hostname": hostname,
                    "port": port,
                    "username": username
                }
            }
        else:
            return {
                "success": False,
                "message": f"连接到 {hostname}:{port} 失败"
            }
    except Exception as e:
        return {
            "success": False,
            "message": f"连接过程中发生错误: {str(e)}"
        }

@mcp.tool
def execute(command: str, hostname: str = None, username: str = None, password: str = None, port: int = 22, key_filename: str = None) -> dict:
    """执行SSH命令
    
    Args:
        command (str): 要在远程服务器上执行的命令
        hostname (str, optional): SSH服务器主机名或IP地址，如果已连接则不需要
        username (str, optional): SSH登录用户名，如果已连接则不需要
        password (str, optional): SSH登录密码，如果使用密钥认证则不需要
        port (int, optional): SSH服务器端口号，默认为22
        key_filename (str, optional): SSH私钥文件路径，如果使用密码认证则不需要
        
    Returns:
        dict: 命令执行结果，包含以下字段：
            - success (bool): 命令是否执行成功
            - command (str): 执行的命令内容
            - output (str): 命令的标准输出内容
            - error (str, optional): 命令的错误输出内容
    """
    try:
        # 如果提供了新的连接信息，先连接
        if hostname and username:
            connection = SSHConnection(
                hostname=hostname,
                port=port,
                username=username,
                password=password,
                key_filename=key_filename
            )
            if not ssh_manager.connect(connection):
                return {
                    "success": False,
                    "message": "SSH连接失败"
                }
        
        # 执行命令
        result = ssh_manager.execute_command(command)
        
        return {
            "success": result.success,
            "command": command,
            "output": result.output,
            "error": result.error
        }
        
    except Exception as e:
        return {
            "success": False,
            "message": f"命令执行过程中发生错误: {str(e)}"
        }

@mcp.tool
def disconnect() -> dict:
    """断开SSH连接
    
    Returns:
        dict: 断开连接结果，包含以下字段：
            - success (bool): 操作是否成功
            - message (str): 操作状态消息
    """
    ssh_manager.disconnect()
    return {
        "success": True,
        "message": "SSH连接已断开"
    }

@mcp.tool
def status() -> dict:
    """检查SSH连接状态
    
    Returns:
        dict: 连接状态信息，包含以下字段：
            - connected (bool): 当前是否已连接
            - message (str): 连接状态描述
    """
    is_connected = ssh_manager.client is not None and ssh_manager.client.get_transport() and ssh_manager.client.get_transport().is_active()
    
    return {
        "connected": is_connected,
        "message": "已连接" if is_connected else "未连接"
    }

@mcp.tool
def add(a: int, b: int) -> int:
    """计算两个数字的和
    
    Args:
        a (int): 第一个数字
        b (int): 第二个数字
        
    Returns:
        int: 两个数字的和
    """
    return a + b

if __name__ == "__main__":
    mcp.run()